package srst.ai;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import org.apache.log4j.Logger;
import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.word2vec.Word2Vec;
import org.deeplearning4j.text.sentenceiterator.BasicLineIterator;
import org.deeplearning4j.text.sentenceiterator.SentenceIterator;
import org.nd4j.linalg.api.ndarray.INDArray;

import srst.util.FilePathUtil;




/**
 * @author YuanYuhu
 *
 */
public class Word2vecUtil {
	
	private static Logger log = Logger.getLogger(Word2vecUtil.class.getName());
	
	public static Word2Vec getWord2Vec(String segmentedFilePath, String wordVectorPath) throws IOException {
		log.info("Load & Vectorize Sentences....");
		// Strip white space before and after for each line
		SentenceIterator iter = new BasicLineIterator(segmentedFilePath);
		// Split on white spaces in the line to get words
		log.info("Building model....");
		Word2Vec word2Vec = new Word2Vec.Builder().minWordFrequency(2).iterations(10).layerSize(100).seed(42).windowSize(5)
				.iterate(iter).build();
		log.info("Fitting Word2Vec model....");
		word2Vec.fit();
		log.info("Writing word vectors to text file....");
		WordVectorSerializer.writeWordVectors(word2Vec, wordVectorPath);
//		String name = "计算机科学与技术";
//		log.info("Closest Words:");
//
//		System.out.println(name + ">>>>>>");
//		Collection<String> lst = word2Vec.wordsNearest(name, 20);
//		System.out.println(lst);
		return word2Vec;
	}
	
//	public static void main(String[] args){
//		
//		try {
//			Word2vecUtil.getWord2Vec(FilePathUtil.SEGMENT_RESUME_PATH,FilePathUtil.WORD_EMBEDDING_VECTOR_PATH);
//		} catch (IOException e) {
//			// TODO Auto-generated catch block
//			e.printStackTrace();
//		}
//	}

}















